package cn.edu.bjtu.test.model.ui;

import java.awt.Component;
import java.awt.GridBagConstraints;
import java.awt.GridBagLayout;
import java.io.File;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.atomic.AtomicInteger;

import javax.swing.JButton;
import javax.swing.JCheckBox;
import javax.swing.JFileChooser;
import javax.swing.JFrame;
import javax.swing.JLabel;
import javax.swing.JPanel;
import javax.swing.JScrollPane;
import javax.swing.JTextArea;

import org.deeplearning4j.api.storage.StatsStorage;
import org.deeplearning4j.nn.conf.ComputationGraphConfiguration;
import org.deeplearning4j.nn.conf.ConvolutionMode;
import org.deeplearning4j.nn.conf.MultiLayerConfiguration;
import org.deeplearning4j.nn.conf.NeuralNetConfiguration;
import org.deeplearning4j.nn.conf.Updater;
import org.deeplearning4j.nn.conf.graph.MergeVertex;
import org.deeplearning4j.nn.conf.layers.ConvolutionLayer;
import org.deeplearning4j.nn.conf.layers.GlobalPoolingLayer;
import org.deeplearning4j.nn.conf.layers.OutputLayer;
import org.deeplearning4j.nn.conf.layers.PoolingType;
import org.deeplearning4j.nn.weights.WeightInit;
import org.deeplearning4j.ui.api.UIServer;
import org.deeplearning4j.ui.storage.InMemoryStatsStorage;
import org.nd4j.linalg.activations.Activation;
import org.nd4j.linalg.lossfunctions.LossFunctions;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import cn.edu.bjtu.configuration.TextCategorizationCNNConfig;
import cn.edu.bjtu.core.CNNNetworkDesignHandler;
import cn.edu.bjtu.core.TextCategorizationManager.ClassificationPair;
import cn.edu.bjtu.model.TextCategorizationCNNModel;
import scala.collection.mutable.StringBuilder;

public class MainUI extends JFrame implements CNNNetworkDesignHandler{
	private Logger logger = LoggerFactory.getLogger(this.getClass());
	private static final long serialVersionUID = -207110178573726271L;
	private JPanel mainJP = null;
	private GridBagLayout gl = null;
	//模型
	private TextCategorizationCNNModel model = null;
	private TextCategorizationCNNConfig config = TextCategorizationCNNConfig.get();
	//控件
	private JTextArea area;
	private JButton loadBtn;
	private JButton classifyBtn;
	private JButton trainButton;
	private JTextArea classifyArea;
	private JButton showConfigBtn;
	private JButton analysisFile;
	private JButton showModel;
	private JButton saveModel;
	
	private JCheckBox openServer ;
	
	UIServer server = null;
	private ThreadFactory _tf = new ThreadFactory() {
		AtomicInteger at = new AtomicInteger(0);
		@Override
		public Thread newThread(Runnable r) {
			Thread  t = new Thread(r);
			t.setName("MainUI-Event-Handler-Pool-"+at.getAndIncrement());
			return t;
		}
	};
	private ExecutorService eventHandler =  Executors.newFixedThreadPool(5,_tf);
	
	//状态
	private volatile boolean isTraining  = false;
	
	public MainUI() {
		super.setDefaultCloseOperation(JFrame.EXIT_ON_CLOSE);
		gl  = new GridBagLayout();
		mainJP = new JPanel(gl);
		setSize(1024, 768);
		setTitle("CNN 测试");
		init();
		setContentPane(mainJP);
		setVisible(true);
		initEventHandler();
		showMessage("!!!点击训练模型之后,配置文件中的指向的模型保存文件将会被覆盖,一个注意保存", true);
	}
	public TextCategorizationCNNConfig getConfig() {
		return config;
	}
	/**
	 * 控件布局函数
	 */
	private void init(){
		//X往右,Y往下
		layout(new JLabel("Model测试"),new GBC(0,0).setGridWidth(5).setGridHeight(1).setFill(GridBagConstraints.BOTH).setWx(1).setWy(1));
		area = new JTextArea("",30,30);
		area.setLineWrap(true);
		JScrollPane js = new JScrollPane(area);
		js.setVerticalScrollBarPolicy(JScrollPane.VERTICAL_SCROLLBAR_AS_NEEDED);
		js.setHorizontalScrollBarPolicy(JScrollPane.HORIZONTAL_SCROLLBAR_NEVER);
		layout(js,new GBC(0, 1).setFill(GridBagConstraints.BOTH).setWx(1).setWy(1).setGridHeight(10).setGridWidth(10));
		classifyArea = layout(new JTextArea("Result goes here"),new GBC(10,1).setGridHeight(10).setGridWidth(10).setFill(GridBagConstraints.BOTH).setWx(1).setWy(1));
		classifyArea.setLineWrap(true);
		
		loadBtn = layout(new JButton("加载模型"),new GBC(0,14));
		showModel = layout(new JButton("查看模型参数"),new GBC(1,14));
		classifyBtn = layout(new JButton("分类"),new GBC(2,14));
		trainButton = layout(new JButton("训练模型"),new GBC(3,14));
		saveModel = layout(new JButton("保存模型"),new GBC(4,14));
		showConfigBtn = layout(new JButton("查看配置"),new GBC(5,14));
		analysisFile = layout(new JButton("打开文件"),new GBC(6,14));
		layout(new JLabel("训练模型时打开UIServer"),new GBC(7,14).setFill(GBC.BOTH).setAnchor(GBC.EAST));
		openServer = layout(new JCheckBox(),new GBC(8,14).setAnchor(GBC.EAST).setFill(GBC.BOTH));
	}
	/**
	 * 为相应的button添加事件监听
	 */
	private void initEventHandler(){
		loadBtn.addActionListener((e)->{
			logger.info("Event");
			eventHandler.submit(()->{
				if(model == null){
					model = TextCategorizationCNNModel.get();
					model.restoreFromFile();
					showMessage("加载完成",true);
				}
			});
		});
		classifyBtn.addActionListener((e)->{
			logger.info("Event");
			String text = area.getText().trim();
			if(text.length()>0){
				eventHandler.execute(()->{
					ClassificationPair[] res =  model.predictDocumentLabelString(text);
					StringBuilder sb = new StringBuilder();
					for(int i = 0; i<res.length;i++){
						sb.append(String.format("%-8s:",res[i].getLabelName()));
						sb.append(String.format("%5f\r\n", res[i].getVal()));
					}
					showMessage(sb.toString(),false);
				});
			}
		});
		trainButton.addActionListener((e)->{
			logger.info("Event");
			eventHandler.submit(()->{
				if(!isTraining){
					//不会有并发问题,但是防止重复训练浪费时间
					isTraining = true;
					trainButton.setEnabled(false);
					if(model == null){
						model = TextCategorizationCNNModel.get();
					}
					if(openServer.isSelected()){
						server = UIServer.getInstance();
						StatsStorage stats = new InMemoryStatsStorage();
						model.config(stats);
						server.attach(stats);
					}else{
						model.config(null);
					}
					model.configNetworkDesignHandler(this);
					model.buildNetworkModel();
					trainButton.setEnabled(true);
					isTraining = false;
				}
				return model;
			});
		});
		showConfigBtn.addActionListener((e)->{
			logger.info("Event");
			new ConfigFrame(MainUI.this);
		});
		analysisFile.addActionListener((e)->{
			logger.info("Event");
			JFileChooser jf = new JFileChooser(new File(config.getDataSetDirOrFile()).getParent());
			int res = jf.showOpenDialog(this);
			if(res == JFileChooser.APPROVE_OPTION){
				new DataFileView(this, jf.getSelectedFile());
			}
			
		});
		showModel.addActionListener((e)->{
			new ShowModelJFrame(this);
		});
		openServer.setSelected(true);
		
		saveModel.addActionListener((e)->{
			JFileChooser jf = new JFileChooser(new File(config.getNetworkSaveFile()).getParent());
			int res = jf.showOpenDialog(this);
			if(res == JFileChooser.APPROVE_OPTION){
				File modelSaveFile = jf.getSelectedFile();
				model.saveCNNModel(modelSaveFile);
			}
		});
		
	}
	
	


	private void showMessage(String msg,boolean left){
		if(left){
			area.setText(msg);
		}else{
			classifyArea.setText(msg);
		}
	
	}
	private <T extends Component> T layout(T comp,GBC gbc){
		mainJP.add(comp,gbc);
		return comp;
	}
	public ComputationGraphConfiguration handle(int embeddingWordVectorLength) {
		int vectorSize = embeddingWordVectorLength;
		int cnnLayerFeatureMaps = 100;
		ComputationGraphConfiguration config = new NeuralNetConfiguration.Builder()
	            .weightInit(WeightInit.RELU)
	            .activation(Activation.LEAKYRELU)
	            .updater(Updater.ADAM)
	            .convolutionMode(ConvolutionMode.Same)      //This is important so we can 'stack' the results later
	            .regularization(true).l2(0.0001)
	            .learningRate(0.01)
	            .graphBuilder()
	            .addInputs("input")
	            .addLayer("cnn3", new ConvolutionLayer.Builder()
	                .kernelSize(3,vectorSize)
	                .stride(1,vectorSize)
	                .nIn(1)
	                .nOut(cnnLayerFeatureMaps)
	                .build(), "input")
	            .addLayer("cnn4", new ConvolutionLayer.Builder()
	                .kernelSize(4,vectorSize)
	                .stride(1,vectorSize)
	                .nIn(1)
	                .nOut(cnnLayerFeatureMaps)
	                .build(), "input")
	            .addLayer("cnn5", new ConvolutionLayer.Builder()
	                .kernelSize(5,vectorSize)
	                .stride(1,vectorSize)
	                .nIn(1)
	                .nOut(cnnLayerFeatureMaps)
	                .build(), "input")
	            .addVertex("merge", new MergeVertex(), "cnn3", "cnn4", "cnn5")
	            .addLayer("globalPool", new GlobalPoolingLayer.Builder()
	                .poolingType(PoolingType.MAX)
	                .build(), "merge")
	            
	            /////////////////////
	            .addLayer("cnn6", new ConvolutionLayer.Builder()
		                .kernelSize(3,vectorSize)
		                .stride(1,vectorSize)
		                .nIn(3*cnnLayerFeatureMaps)
		                .nOut(cnnLayerFeatureMaps)
		                .build(), "merge")
		            .addLayer("cnn7", new ConvolutionLayer.Builder()
		                .kernelSize(4,vectorSize)
		                .stride(1,vectorSize)
		                .nIn(3*cnnLayerFeatureMaps)
		                .nOut(cnnLayerFeatureMaps)
		                .build(), "merge")
		            .addLayer("cnn8", new ConvolutionLayer.Builder()
		                .kernelSize(5,vectorSize)
		                .stride(1,vectorSize)
		                .nIn(3*cnnLayerFeatureMaps)
		                .nOut(cnnLayerFeatureMaps)
		                .build(), "merge")
		            .addVertex("merge2", new MergeVertex(), "cnn6", "cnn7", "cnn8")
		            .addLayer("globalPool2", new GlobalPoolingLayer.Builder()
		                .poolingType(PoolingType.MAX)
		                .build(), "merge2")
		            /////////////////////////
	            .addLayer("out", new OutputLayer.Builder()
	                .lossFunction(LossFunctions.LossFunction.MCXENT)
	                .activation(Activation.SOFTMAX)
	                .nIn(/*3*/cnnLayerFeatureMaps)
	                .nOut(13)    
	                .build(), "globalPool2")
	            .setOutputs("out")
	            .build();
		return config;
	}
	/**
	 * @Param
	 * @Author zhangzhidong
	 */
	
	@Override
	public ComputationGraphConfiguration handleCGC(int embeddingWordVectorLength) {
		// TODO Auto-generated method stub
		return null;
	}
	/**
	 * @Param
	 * @Author zhangzhidong
	 */
	
	@Override
	public MultiLayerConfiguration handleMLN(int embeddingWordVectorLength) {
		// TODO Auto-generated method stub
		return null;
	}
}